package cn.cas.mango.service;


import cn.cas.mango.util.Pair;
import cn.cas.mango.util.TimeUtil;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.bulk.BulkOperation;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.web.client.RestTemplateBuilder;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
@Slf4j
public class RequestService {

  @Value("${model.url}")
  private String modelUrl;

  @Value("${model.field}")
  private String field;

  @Autowired
  private ElasticsearchClient client;

  private List<List<Double>> records = new ArrayList<>();
  public String get(List<String> contents) {
    contents = contents.stream().map(x->x == null ? "": x).collect(Collectors.toList());
    JSONObject request = new JSONObject();
    request.put("category", "");
    request.put("newsList", contents.stream().map(String::trim).collect(Collectors.toList()));
    RestTemplateBuilder builder = new RestTemplateBuilder();
    RestTemplate restTemplate = builder.build();
    String url = modelUrl;
    TimeUtil start = new TimeUtil();
    ResponseEntity<JSONObject> responseEntity = restTemplate.postForEntity(url, request,
        JSONObject.class);
    JSONObject body = responseEntity.getBody();
    Double cost = start.getCost();
    int length = contents.stream().map(x->x.length()).reduce((a,b)->a+b).get();
    List<Double> list = Arrays.asList((double) length, (double) contents.size(), cost);
    records.add(list);

    log.info("length: {}, size: {}, cost: {}s", length, contents.size(), cost);
    assert body != null;
    return body.getString(field);
  }

  // delete
  public void getAllCost() {
    log.info("model length {}, total: {}, size: {}, cost: {}", records.stream().map(x->x.get(0)).reduce((a,b)->a+b).get().intValue()
        , records.size(),
        records.stream().map(x->x.get(1)).reduce((a,b)->a+b).get().intValue(),
        String.format("%.3fs", records.stream().map(x->x.get(2)).reduce((a,b)->a+b).get()));
    records.clear();
  }

  public  <T> List<T> getLarge(SearchRequest.Builder builder1, SearchRequest.Builder builder2, Class<T> tclass) {
    long total;
    List<T> ret = new ArrayList<>();
    try {
      SearchResponse<Map> search = client.search(builder1.size(0).trackTotalHits(v->v.enabled(true)).build(), Map.class);
      total = search.hits().total().value();
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
    int batch = 2000;
    for (int i = 0; i <= (total - 1) / batch; i++) {
      int finalI = i;

      SearchResponse<T> search = null;
      SearchRequest.Builder builder = new SearchRequest.Builder();
      SearchRequest build = builder2.from(finalI * batch).size(batch).build();
      builder.index(build.index());
      builder.query(build.query());
      builder2 = builder;
      try {
        search = client.search(build, tclass);
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
      List<T> collect = search.hits().hits().stream().map(x -> x.source()).collect(Collectors.toList());
      ret.addAll(collect);
    }
    return ret;
  }

  private List<List<BulkOperation>> splitBatch(List<BulkOperation> clusterOperations) {
    int splitSize = 1000;
    int size = clusterOperations.size();
    int count = (size - 1) / splitSize + 1;
    List<List<BulkOperation>> ret = new ArrayList<>();
    for (int i = 0; i < count; i++) {
      List<BulkOperation> bulkOperations = clusterOperations.subList(i * splitSize, Math.min(size, i * splitSize + splitSize));
      ret.add(bulkOperations);
    }
    return ret;
  }

  public void bulkWrite(List<BulkOperation> bulk) {
    splitBatch(bulk).forEach(x-> {
      try {
        client.bulk(y->y.operations(x));
      } catch (IOException e) {
        throw new RuntimeException(e);
      }
    });
  }
}
